#!/usr/bin/env python3
"""
网络连接诊断脚本
检查服务器是否能访问 HuggingFace Hub
"""

import subprocess
import requests
import os
import sys

def check_dns(domain):
    """检查 DNS 解析"""
    print(f"\n[1] 检查 DNS 解析: {domain}")
    try:
        import socket
        ip = socket.gethostbyname(domain)
        print(f"✓ DNS 解析成功: {domain} -> {ip}")
        return True
    except Exception as e:
        print(f"✗ DNS 解析失败: {e}")
        return False

def check_ping(domain):
    """检查 ping 连通性"""
    print(f"\n[2] 检查 Ping 连通性: {domain}")
    try:
        # Windows 和 Linux 的 ping 命令略有不同
        param = "-n" if os.name == "nt" else "-c"
        result = subprocess.run(
            ["ping", param, "1", domain],
            capture_output=True,
            text=True,
            timeout=10
        )
        if result.returncode == 0:
            print(f"✓ Ping 成功")
            return True
        else:
            print(f"✗ Ping 失败")
            return False
    except Exception as e:
        print(f"✗ Ping 命令执行失败: {e}")
        return False

def check_https(url):
    """检查 HTTPS 连接"""
    print(f"\n[3] 检查 HTTPS 连接: {url}")
    try:
        response = requests.get(url, timeout=10)
        print(f"✓ HTTPS 连接成功, 状态码: {response.status_code}")
        return True
    except requests.exceptions.Timeout:
        print(f"✗ 连接超时")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"✗ 连接错误: {e}")
        return False
    except Exception as e:
        print(f"✗ 未知错误: {e}")
        return False

def check_huggingface_api():
    """检查 HuggingFace API"""
    print(f"\n[4] 检查 HuggingFace API")
    api_url = "https://huggingface.co/api/models/stabilityai/stable-video-diffusion-img2vid-xt"
    try:
        response = requests.get(api_url, timeout=10)
        if response.status_code == 200:
            print(f"✓ HuggingFace API 可访问")
            return True
        else:
            print(f"✗ API 返回错误状态码: {response.status_code}")
            return False
    except Exception as e:
        print(f"✗ 无法访问 HuggingFace API: {e}")
        return False

def check_proxy():
    """检查代理设置"""
    print(f"\n[5] 检查代理设置")
    http_proxy = os.environ.get('HTTP_PROXY') or os.environ.get('http_proxy')
    https_proxy = os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
    
    if http_proxy:
        print(f"  HTTP_PROXY: {http_proxy}")
    else:
        print(f"  HTTP_PROXY: 未设置")
    
    if https_proxy:
        print(f"  HTTPS_PROXY: {https_proxy}")
    else:
        print(f"  HTTPS_PROXY: 未设置")
    
    return http_proxy or https_proxy

def suggest_solutions(results):
    """根据检查结果提供解决方案"""
    print("\n" + "="*50)
    print("诊断结果与建议:")
    print("="*50)
    
    if not results['dns']:
        print("\n⚠ DNS 解析失败")
        print("建议:")
        print("  1. 检查 /etc/resolv.conf (Linux) 或网络设置")
        print("  2. 尝试使用公共 DNS: 8.8.8.8 或 1.1.1.1")
    
    if not results['ping']:
        print("\n⚠ Ping 不通")
        print("建议:")
        print("  1. 检查防火墙设置")
        print("  2. 检查路由配置")
    
    if not results['https']:
        print("\n⚠ HTTPS 连接失败")
        print("建议:")
        print("  1. 检查防火墙是否允许 443 端口")
        print("  2. 检查是否需要配置代理")
    
    if not results['api']:
        print("\n⚠ HuggingFace API 无法访问")
        print("建议:")
        print("  1. 考虑使用镜像站点")
        print("  2. 配置代理服务器")
        print("  3. 使用手动下载方式")
    
    if not results['proxy']:
        print("\n💡 未检测到代理设置")
        print("如果您的网络需要代理，请设置:")
        print("  export HTTP_PROXY=http://your-proxy:port")
        print("  export HTTPS_PROXY=http://your-proxy:port")

def main():
    print("="*50)
    print("HuggingFace 网络连接诊断工具")
    print("="*50)
    
    results = {
        'dns': check_dns('huggingface.co'),
        'ping': check_ping('huggingface.co'),
        'https': check_https('https://huggingface.co'),
        'api': check_huggingface_api(),
        'proxy': check_proxy()
    }
    
    suggest_solutions(results)
    
    # 如果所有检查都失败，返回错误码
    if not any([results['dns'], results['https'], results['api']]):
        print("\n❌ 网络连接存在严重问题，无法访问 HuggingFace")
        sys.exit(1)
    else:
        print("\n✅ 部分网络功能正常")

if __name__ == "__main__":
    main()
